fFitAndPlot <- function (dat, 
                         model1, 
                         model2, 
                         model3, 
                         model4, 
                         model5)
  
{  
  
  fit1  <- fFitModel(model1, dat)
  fit2  <- fFitModel(model2, dat)
  fit3  <- fFitModel(model3, dat)
  fit4  <- fFitModel(model4, dat)
  fit5  <- fFitModel(model5, dat)
  
  loo1 <- loo::loo(fit1$fit, save_psis = TRUE, 'logLikelihood')
  loo2 <- loo::loo(fit2$fit, save_psis = TRUE, 'logLikelihood')
  loo3 <- loo::loo(fit3$fit, save_psis = TRUE, 'logLikelihood')
  loo4 <- loo::loo(fit4$fit, save_psis = TRUE, 'logLikelihood')
  loo5 <- loo::loo(fit5$fit, save_psis = TRUE, 'logLikelihood')
  
  comparison <- data.frame(round(loo::loo_compare(x = list(loo1, loo2, loo3, loo4, loo5)), 3)) [, 1:2]
  
  pp1 <- fPlotModel(fit1, dat, 'model 1', 'constant foi')
  pp2 <- fPlotModel(fit2, dat, 'model 2', 'continuous')
  pp3 <- fPlotModel(fit3, dat, 'model 3', 'Student t ')
  pp4 <- fPlotModel(fit4, dat, 'model 4', 'contin. unsmoothed')
  pp5 <- fPlotModel(fit5, dat, 'model 5', 'laplace')
  
  ppC <-  
      ggplot(data.frame(x= 0:1, y = 0:1), aes(x, y)) +
      geom_blank() + ylab('') + xlab ('') + theme_void() +
      annotation_custom(tableGrob(comparison)) 
  
  res_survey <- list(dat  = dat,
                     fit1 = fit1,
                     fit2 = fit2,
                     fit3 = fit3,
                     fit4 = fit4,
                     fit5 = fit5,
                     comparison = comparison,
                     pp1  = pp1,
                     pp2  = pp2,
                     pp3  = pp3,
                     pp4  = pp4,
                     pp5  = pp5,
                     ppC  = ppC
  )
  
  return(res_survey)
  
}



# ---- Datasets
dat1 <- readRDS('data/Panama2012.RDS') %>% mutate(id = stringr::str_sub(dataset_id, 1, 3))
dat2 <- readRDS('data/Panama2017.RDS') %>% mutate(id = stringr::str_sub(dataset_id, 1, 3))

dat0   <- rbind(dat1, dat2) %>% 
  mutate(counts = pos,
         survey = paste(virus,tsur, id)) %>%
  arrange(survey)
datasets <- unique(dat0$survey)



for (s in datasets) 
{
  dat <- filter(dat0, survey == s) %>% arrange(age_mean_f)
  res <- fFitAndPlot(dat, model1, model2, model3, model4, model5)
  saveRDS(res, paste0('res/', s, '.RDS' ))
  gridExtra::grid.arrange(res$pp1, res$pp2, res$pp3, 
                          res$pp4, res$pp5, res$ppC,
                          nrow = 1)
  grid.rect(width = 1, height = 1, gp = gpar(lwd = 2, col = "blue", fill = NA))
  rm(dat, res)
}